import zipfile
import rarfile
import tkinter
from tkinter import filedialog
import os
import time

def selectFile():
    root = tkinter.Tk()
    root.withdraw()
    # root.mainloop()

    file = filedialog.askopenfilename()
    # fileDir = os.path.dirname(file)
    fileDir = os.path.splitext(file)[0]
    return (file, fileDir)

def checkFile(file):
    if not file:
        print('没有文件被选择！')
        return False

    if os.path.splitext(file)[1] not in ['.zip','.rar']:
        print('选择的不是有效的压缩文件！')
        return False

    return True

def extract(zFile, toPath, password=None):
    try:
        if isinstance(zFile,zipfile.ZipFile):
            zFile.extractall(path=toPath, pwd=str(password).encode('utf-8'))
        elif isinstance(zFile,rarfile.RarFile):
            # rarfile.UNRAR_TOOL = 'unrar.exe'
            zFile.extractall(path=toPath, pwd=str(password))
        print('======== 解压文件如下 ========')
        for filename in zFile.namelist():
            print(os.path.join(toPath,filename))
        return password
    except Exception as e:
        print(e)
        pass
        return

def isEncrypted(zFile):
    if hasattr(zFile,'needs_password'):
        # rar文件
        if zFile.needs_password():
            return True
    elif hasattr(zFile,'infolist'):
        #zip文件
        for i in zFile.infolist():
            if (i.flag_bits & 0x01):
                return True
    return False

def getZFile(file):
    if '.zip' in os.path.splitext(file):
        zFile = zipfile.ZipFile(file)
    elif '.rar' in os.path.splitext(file):
        zFile = rarfile.RarFile(file)
    return zFile

def preCrack(zFile):

    passfile = open('./pass.txt')

    if not isEncrypted(zFile):
        print('该文件未加密，无需破解！')
        return
    return zFile


def crack(zFile,toPath):
    try:
        for nums in range(0, 6):
            for password in range(pow(10, nums), pow(10, nums + 1)):
                # for line in passfile.readlines():
                #         password = line.strip('\n')
                # password = password.replace('\n','')
                print('使用密码 %s' % password)
                guess = extract(zFile, toPath, password)
                if guess:
                    print('破解成功，密码是：' + str(password))
                    return

        print('密码过于复杂，破解失败！')

    except Exception as e:
        print(e)
    finally:
        # passfile.close()
        zFile.close()

def main():
    (file, toPath) = selectFile()

    if not checkFile(file):
        exit(0)

    zFile = getZFile(file)
    needCrack = preCrack(zFile)

    start = time.clock()
    if needCrack:
        # crack(*file)
        crack(zFile, toPath)
    else:
        extract(zFile, toPath)
    end = time.clock()
    print('========= 总用时：%.2fs ==========' % (end - start))


if __name__ == '__main__':
    main()